use crate::models::{Attrs, Message, Role};
use anyhow::{anyhow, Context, Result};
use serde_json::Value;
use std::collections::HashMap;
use std::fs;
use std::io::prelude::*;
use std::path::Path;

const MESSAGES_FILE: &str = "messages";

pub fn prepare_messages(
    system_message: &Message,
    question: &str,
    kb_number: &Option<usize>,
) -> (Message, String) {
    if let Some(kb_message) = super::get_kb_message(kb_number) {
        let kb_system_message = Message {
            role: Role::System,
            content: format!("{}\n\n{}", &system_message.content, &kb_message),
            attrs: None,
        };
        (kb_system_message, question.to_string())
    } else {
        (system_message.clone(), question.to_string())
    }
}

pub fn create_user_message(question: &str) -> Message {
    Message {
        role: Role::User,
        content: question.to_string(),
        attrs: Some(Attrs { index: 1 }),
    }
}

pub fn create_assistant_message(response: &str) -> Message {
    Message {
        role: Role::Assistant,
        content: response.to_string(),
        attrs: Some(Attrs { index: 1 }),
    }
}

pub fn load_messages(working_path: &Path) -> Result<Vec<Message>> {
    let messages_file_path = working_path.join(MESSAGES_FILE);

    // 检查文件是否存在
    if !messages_file_path.exists() {
        return Ok(Vec::new());
    }

    // 读取文件内容
    let messages_str = fs::read_to_string(&messages_file_path)?;

    // 解析文件内容
    parse_messages(&messages_str)
}

fn parse_attrs(line_str: &str) -> HashMap<String, String> {
    // 首先找到 `attrs="..."` 的部分
    let find_target: &str = " attrs=\"";
    let target_len: usize = find_target.len();
    if let Some(attrs_start) = line_str.find(find_target) {
        let attrs_end =
            line_str[attrs_start + target_len..].find('\"').unwrap_or(0) + attrs_start + target_len;
        let attrs_str = &line_str[attrs_start + target_len..attrs_end];

        // 解析键值对
        attrs_str
            .split(';')
            .filter_map(|attr| {
                let mut parts = attr.splitn(2, ':');
                if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
                    Some((key.trim().to_string(), value.trim().to_string()))
                } else {
                    None
                }
            })
            .collect()
    } else {
        HashMap::new() // 如果没有找到 `attrs` 属性，返回空 HashMap
    }
}

// 从attrs字符串中提取index属性，如果没有则返回0
fn extract_index(line_str: &str) -> usize {
    let attrs = parse_attrs(line_str);
    attrs.get("index").and_then(|s| s.parse().ok()).unwrap_or(0)
}

fn parse_messages(messages_str: &str) -> Result<Vec<Message>> {
    let mut messages = Vec::new();
    let lines: Vec<&str> = messages_str.lines().collect();
    let mut user_message_start = false;
    let mut ai_message_start = false;
    let mut has_initial_user_message = false;
    let mut current_message = String::new();
    let mut previous_role: Option<Role> = None;
    let mut current_index = 0;

    for (index, line) in lines.iter().enumerate() {
        if line.trim().starts_with("<! user-message-start") {
            if previous_role == Some(Role::User) {
                return Err(anyhow!(
                    "Invalid message order(user_message_start:{})",
                    index
                ));
            }
            user_message_start = true;
            ai_message_start = false;
            current_message.clear();
            current_index = extract_index(line.trim());
        } else if line.trim() == "<! user-message-end>" {
            if user_message_start {
                messages.push(Message {
                    role: Role::User,
                    content: current_message.clone(),
                    attrs: if current_index > 0 {
                        Some(Attrs {
                            index: current_index,
                        })
                    } else {
                        None
                    },
                });
                user_message_start = false;
                has_initial_user_message = true;
                previous_role = Some(Role::User);
            } else {
                return Err(anyhow!("Missing user-message-start"));
            }
        } else if line.trim().starts_with("<! ai-message-start") {
            if previous_role == Some(Role::User) || has_initial_user_message {
                user_message_start = false;
                ai_message_start = true;
                current_message.clear();
                current_index = extract_index(line.trim());
            } else {
                return Err(anyhow!("Invalid message order(ai_message_start:{})", index));
            }
        } else if line.trim() == "<! ai-message-end>" {
            if ai_message_start {
                messages.push(Message {
                    role: Role::Assistant,
                    content: current_message.clone(),
                    attrs: if current_index > 0 {
                        Some(Attrs {
                            index: current_index,
                        })
                    } else {
                        None
                    },
                });
                ai_message_start = false;
                previous_role = Some(Role::Assistant);
            } else {
                return Err(anyhow!("Missing ai-message-start"));
            }
        } else if user_message_start || ai_message_start {
            current_message.push_str(line);
            current_message.push('\n');
        }
    }

    if messages.is_empty() && !current_message.is_empty() {
        return Err(anyhow!("Missing end"));
    }

    Ok(messages)
}

pub fn sort_messages(messages: &[Message]) -> Vec<Message> {
    let len = messages.len();
    messages
        .iter()
        .enumerate()
        .map(|(i, message)| {
            let mut new_message = message.clone();
            if let Some(attrs) = &mut new_message.attrs {
                attrs.index = len - i;
            } else {
                new_message.attrs = Some(Attrs { index: len - i });
            }
            new_message
        })
        .collect()
}

pub fn save_messages(messages: &[Message]) -> Result<()> {
    // 将 messages 序列化为字符串
    let messages_str = serialize_messages(messages)?;

    // 打开文件并写入内容
    let mut file = fs::OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true) // 清空文件内容
        .open(MESSAGES_FILE)
        .context("Failed to open file")?;

    file.write_all(messages_str.as_bytes())
        .context("Failed to write to file")?;

    Ok(())
}

fn serialize_messages(messages: &[Message]) -> Result<String> {
    let mut result = String::new();

    for message in messages {
        // Skip system and tool messages as they are not part of the conversation history
        if matches!(message.role, Role::System | Role::Tool) {
            continue;
        }

        let attrs_str = if let Some(attrs) = &message.attrs {
            if attrs.index > 0 {
                format!(" attrs=\"index:{}\"", attrs.index)
            } else {
                String::new()
            }
        } else {
            String::new()
        };

        match message.role {
            Role::User => {
                result.push_str(&format!("<! user-message-start{}>\n", attrs_str));
                result.push_str(&message.content);
                result.push_str("\n<! user-message-end>\n");
            }
            Role::Assistant => {
                result.push_str(&format!("<! ai-message-start{}>\n", attrs_str));
                result.push_str(&message.content);
                result.push_str("\n<! ai-message-end>\n");
            }
            Role::System | Role::Tool => {} // Already handled above
        }
    }

    Ok(result)
}

pub fn get_latest_messages(messages: &[Message], count: usize) -> Vec<Message> {
    if count >= messages.len() {
        messages.to_vec()
    } else {
        messages
            .iter()
            .skip(messages.len() - count)
            .cloned()
            .collect()
    }
}

#[cfg(test)]
mod tests_parse_messages {
    use super::*;
    use crate::models::{Attrs, Message};

    #[test]
    fn test_parse_attrs() {
        let attrs_str = r#"<! user-message attrs="index:1;other:value""#;
        let attrs = parse_attrs(attrs_str);
        assert_eq!(attrs.get("index"), Some(&"1".to_string()));
        assert_eq!(attrs.get("other"), Some(&"value".to_string()));
    }

    #[test]
    fn test_extract_index() {
        let attrs_str = r#"<! user-message attrs="index:1;other:value""#;
        assert_eq!(extract_index(attrs_str), 1);
        assert_eq!(extract_index("other:value"), 0);
        assert_eq!(extract_index(""), 0);
    }

    #[test]
    fn test_parse_messages() {
        let messages_str = r#"
<! user-message-start attrs="index:1">
Hello, AI!
<! user-message-end>
<! ai-message-start attrs="index:2">
Hello, User!
<! ai-message-end>
"#;
        let messages = parse_messages(messages_str).unwrap();
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].role, Role::User);
        assert_eq!(messages[0].content, "Hello, AI!".to_string());
        println!("message:0:{:?}", messages[0]);
        assert_eq!(messages[0].attrs.clone().unwrap().index, 1);
        assert_eq!(messages[1].role, Role::Assistant);
        assert_eq!(messages[1].content, "Hello, User!".to_string());
        assert_eq!(messages[1].attrs.as_ref().unwrap().index, 2);
    }

    #[test]
    fn test_parse_messages_no_attrs() {
        let messages_str = r#"
<! user-message-start>
Hello, AI!
<! user-message-end>
<! ai-message-start>
Hello, User!
<! ai-message-end>
"#;
        let messages = parse_messages(messages_str).unwrap();
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].role, Role::User);
        assert_eq!(messages[0].content, "Hello, AI!".to_string());
        assert!(messages[0].attrs.is_none());
        assert_eq!(messages[1].role, Role::Assistant);
        assert_eq!(messages[1].content, "Hello, User!".to_string());
        assert!(messages[1].attrs.is_none());
    }

    #[test]
    fn test_parse_messages_invalid_order() {
        let messages_str = r#"
<! user-message-start>
Hello, AI!
<! user-message-end>
<! user-message-start>
Hello again, AI!
<! user-message-end>
"#;
        let result = parse_messages(messages_str);
        assert!(result.is_err());
    }
}
#[cfg(test)]
mod tests_get_latest_messages {
    use super::*;

    #[test]
    fn test_get_latest_messages() {
        let messages = vec![
            Message {
                role: Role::User,
                content: "Hello".to_string(),
                attrs: Some(Attrs { index: 3 }),
            },
            Message {
                role: Role::Assistant,
                content: "World".to_string(),
                attrs: Some(Attrs { index: 2 }),
            },
            Message {
                role: Role::User,
                content: "Rust".to_string(),
                attrs: Some(Attrs { index: 1 }),
            },
        ];

        // Case 1: count is less than the number of messages
        let result = get_latest_messages(&messages, 2);
        assert_eq!(
            result,
            vec![
                Message {
                    role: Role::Assistant,
                    content: "World".to_string(),
                    attrs: Some(Attrs { index: 2 })
                },
                Message {
                    role: Role::User,
                    content: "Rust".to_string(),
                    attrs: Some(Attrs { index: 1 })
                },
            ]
        );

        // Case 2: count is equal to the number of messages
        let result = get_latest_messages(&messages, 3);
        assert_eq!(result, messages);

        // Case 3: count is greater than the number of messages
        let result = get_latest_messages(&messages, 5);
        assert_eq!(result, messages);

        // Case 4: count is zero
        let result = get_latest_messages(&messages, 0);
        assert_eq!(result, vec![]);
    }
}

#[cfg(test)]
mod tests_serialize_message {
    use super::*;

    #[test]
    fn test_serialize_messages_with_attrs() {
        let messages = vec![
            Message {
                role: Role::User,
                content: "Hello, world!".to_string(),
                attrs: Some(Attrs { index: 1 }),
            },
            Message {
                role: Role::Assistant,
                content: "Hi there!".to_string(),
                attrs: Some(Attrs { index: 2 }),
            },
        ];

        let expected_output = "<! user-message-start attrs=\"index:1\">\nHello, world!\n<! user-message-end>\n<! ai-message-start attrs=\"index:2\">\nHi there!\n<! ai-message-end>\n";
        let output = serialize_messages(&messages).unwrap();
        assert_eq!(output, expected_output);
    }

    #[test]
    fn test_serialize_messages_without_attrs() {
        let messages = vec![
            Message {
                role: Role::User,
                content: "Hello, world!".to_string(),
                attrs: None,
            },
            Message {
                role: Role::Assistant,
                content: "Hi there!".to_string(),
                attrs: None,
            },
        ];

        let expected_output = "<! user-message-start>\nHello, world!\n<! user-message-end>\n<! ai-message-start>\nHi there!\n<! ai-message-end>\n";
        let output = serialize_messages(&messages).unwrap();
        assert_eq!(output, expected_output);
    }

    #[test]
    fn test_serialize_messages_with_zero_index() {
        let messages = vec![
            Message {
                role: Role::User,
                content: "Hello, world!".to_string(),
                attrs: Some(Attrs { index: 0 }),
            },
            Message {
                role: Role::Assistant,
                content: "Hi there!".to_string(),
                attrs: Some(Attrs { index: 0 }),
            },
        ];

        let expected_output = "<! user-message-start>\nHello, world!\n<! user-message-end>\n<! ai-message-start>\nHi there!\n<! ai-message-end>\n";
        let output = serialize_messages(&messages).unwrap();
        assert_eq!(output, expected_output);
    }

    #[test]
    fn test_serialize_messages_with_non_string_content() {
        let messages = vec![
            Message {
                role: Role::User,
                content: true.to_string(),
                attrs: None,
            },
            Message {
                role: Role::Assistant,
                content: 42.to_string(),
                attrs: None,
            },
        ];

        let expected_output = "<! user-message-start>\ntrue\n<! user-message-end>\n<! ai-message-start>\n42\n<! ai-message-end>\n";
        let output = serialize_messages(&messages).unwrap();
        assert_eq!(output, expected_output);
    }

    #[test]
    fn test_serialize_messages_skip_system_tool() {
        let messages = vec![
            Message {
                role: Role::System,
                content: "System message".to_string(),
                attrs: None,
            },
            Message {
                role: Role::User,
                content: "Hello".to_string(),
                attrs: None,
            },
            Message {
                role: Role::Tool,
                content: "Tool message".to_string(),
                attrs: None,
            },
            Message {
                role: Role::Assistant,
                content: "Hi".to_string(),
                attrs: None,
            },
        ];

        let expected_output = "<! user-message-start>\nHello\n<! user-message-end>\n<! ai-message-start>\nHi\n<! ai-message-end>\n";
        let output = serialize_messages(&messages).unwrap();
        assert_eq!(output, expected_output);
    }
}
